# -*- coding: utf-8 -*-
"""
core/代码生成/generator.py

智能代码生成与变更建议模块。
负责构建提示词，调用LLM（如Claude-3.7）生成代码或修改建议，并解析其输出。
"""

import os
import json
import re
from core.utils import logger, read_file_content

# --- LLM API 配置 (模拟) ---
# 实际项目中会从 config.ini 或环境变量读取
CLAUDE_API_KEY = "YOUR_CLAUDE_API_KEY_HERE" # 预留，用户配置
CLAUDE_API_ENDPOINT = "https://api.anthropic.com/v1/messages" # 示例，以Anthropic Claude为例
CLAUDE_MODEL_NAME = "claude-3-opus-20240229" # 或用户指定的 claude-3.7 (如果可用)

class CodeGenerator:
    """
    代码生成器类，封装了与LLM交互以生成或修改代码的逻辑。
    """
    def __init__(self, prompt_template_path):
        """
        初始化代码生成器。

        参数:
            prompt_template_path (str): LLM提示词模板文件路径。
        """
        self.prompt_template = read_file_content(prompt_template_path)
        if not self.prompt_template:
            logger.error(f"未能加载提示词模板: {prompt_template_path}")
            raise ValueError(f"提示词模板加载失败: {prompt_template_path}")
        logger.info(f"代码生成提示词模板已加载: {prompt_template_path}")

    def _construct_llm_prompt(self, requirement_item, rag_context):
        """
        根据需求差异项和RAG检索结果，填充提示词模板。

        参数:
            requirement_item (dict): 单个需求差异项。
            rag_context (dict): RAG检索模块返回的上下文信息，
                               包含 {"related_docs": [...], "related_code": [...]}

        返回:
            str: 填充完毕的完整提示词。
        """
        prompt = self.prompt_template

        # 1. 填充任务类型
        task_type_marker = f"- [x] {requirement_item.get("type")}"
        if requirement_item.get("type") == "新增":
            prompt = prompt.replace("- [ ] 新增功能 (Generate new C code)", task_type_marker)
        elif requirement_item.get("type") == "修改":
            prompt = prompt.replace("- [ ] 修改功能 (Suggest modifications to existing C code)", task_type_marker)
        elif requirement_item.get("type") == "移除":
            prompt = prompt.replace("- [ ] 移除功能 (Suggest C code for removal/commenting out)", task_type_marker)
        
        # 2. 填充新需求描述
        new_desc_placeholder = "[此处插入新需求描述]"
        new_description = requirement_item.get("new_description", 
                             requirement_item.get("feature_name", "") + ": " + requirement_item.get("details", ""))
        prompt = prompt.replace(new_desc_placeholder, new_description.strip())

        # 3. 填充相关旧需求描述
        old_desc_placeholder = "[此处插入相关旧需求描述]"
        old_description = requirement_item.get("old_description", "无相关旧需求信息。")
        prompt = prompt.replace(old_desc_placeholder, old_description.strip())

        # 4. 填充RAG文档检索结果
        doc_results_str = ""
        if rag_context.get("related_docs"):
            for i, doc in enumerate(rag_context["related_docs"]):
                doc_results_str += f"--- 结果 #{i+1} ---\n"
                doc_results_str += f"标题：{doc.get("title", "N/A")}\n"
                doc_results_str += f"距离：{doc.get("distance", "N/A")}\n"
                doc_results_str += f"内容：\n{doc.get("content", "N/A")}\n---\n\n"
        else:
            doc_results_str = "未从《需求联系代码的拆解文档》中检索到相关内容。\n"
        # 定位模板中RAG文档部分并替换 (假设模板中有明确的标记)
        # 这是一个简化的替换，实际模板中可能需要更精确的占位符
        prompt = re.sub(r"(以下是从《需求联系代码的拆解文档》中检索到的.*?)(--- 结果 #1 ---.*?---(\n\n--- 结果 #\d ---.*?---)*)?", 
                        f"\1\n{doc_results_str.strip()}", prompt, flags=re.DOTALL)

        # 5. 填充RAG代码检索结果
        code_results_str = ""
        if rag_context.get("related_code"):
            for i, code_frag in enumerate(rag_context["related_code"]):
                code_results_str += f"**代码片段 {i+1}:**\n"
                code_results_str += f"文件名: `{code_frag.get("file_path", "N/A")}`\n"
                code_results_str += f"函数名/代码块描述: `{code_frag.get("name", "N/A")} ({code_frag.get("type", "N/A")})`\n"
                code_results_str += f"起止行号: `L{code_frag.get("start_line", "?")}-L{code_frag.get("end_line", "?")}`\n"
                code_results_str += f"代码内容:\n```c\n{code_frag.get("code", "// 未找到代码")}\n```\n\n"
        else:
            code_results_str = "未从旧项目代码库中检索到相关代码片段。\n"
        # 定位模板中RAG代码部分并替换
        prompt = re.sub(r"(以下是从旧项目 `TM56F1522-V1.0-OLD` 代码库中检索到的.*?)(代码片段 1:.*?```(\n\n代码片段 \d:.*?```)*)?",
                        f"\1\n{code_results_str.strip()}", prompt, flags=re.DOTALL)
        
        return prompt

    def _call_llm_api(self, prompt_text):
        """
        模拟调用LLM API (例如 Claude-3.7)。
        实际实现中会使用 `requests` 库。

        参数:
            prompt_text (str): 发送给LLM的完整提示词。

        返回:
            str: LLM返回的原始文本响应。如果失败则返回None。
        """
        logger.info("开始调用LLM API进行代码生成/修改建议...")
        logger.debug(f"发送给LLM的最终提示词 (部分预览):\n{prompt_text[:1000]}\n...")

        # ---- 模拟LLM API调用 ----
        # headers = {
        #     "x-api-key": CLAUDE_API_KEY,
        #     "anthropic-version": "2023-06-01",
        #     "content-type": "application/json"
        # }
        # data = {
        #     "model": CLAUDE_MODEL_NAME,
        #     "max_tokens": 4000, # 根据需要调整
        #     "messages": [{"role": "user", "content": prompt_text}]
        # }
        # try:
        #     response = requests.post(CLAUDE_API_ENDPOINT, headers=headers, json=data, timeout=120)
        #     response.raise_for_status() # 如果HTTP错误则抛出异常
        #     response_data = response.json()
        #     llm_output = response_data.get("content", [{}])[0].get("text", "")
        #     logger.info("LLM API调用成功，收到响应。")
        #     return llm_output
        # except requests.exceptions.RequestException as e:
        #     logger.error(f"LLM API调用失败: {e}")
        #     return None
        # except Exception as e:
        #     logger.error(f"处理LLM响应时发生未知错误: {e}")
        #     return None
        
        # ---- 模拟响应 ----
        # 这个模拟响应需要根据提示词模板中定义的输出格式来构造
        simulated_response = """
        建议的 `.h` 文件修改/内容 (如果适用):
        文件名: `key.h`
        ```c
        // key.h
        #ifndef KEY_H
        #define KEY_H

        // 新增童锁状态枚举和函数声明
        typedef enum {
            CHILD_LOCK_OFF = 0,
            CHILD_LOCK_ON = 1
        } ChildLockState_E;

        void KEY_Check_ChildLock_Activation(void);
        ChildLockState_E KEY_Get_ChildLock_State(void);

        #endif // KEY_H
        ```

        建议的 `.c` 文件修改/内容:
        文件名: `key.c`
        ```c
        // key.c
        #include "key.h"
        #include "led.h" // 假设童锁状态会控制LED

        static ChildLockState_E g_child_lock_state = CHILD_LOCK_OFF;
        static unsigned int child_lock_press_time = 0;

        // 新增函数：获取童锁状态
        ChildLockState_E KEY_Get_ChildLock_State(void) {
            return g_child_lock_state;
        }

        // 新增函数：检测童锁激活 (应在按键扫描任务中定期调用)
        // 假设【模式键】+【加键】为组合键 KEY_MODE 和 KEY_PLUS
        void KEY_Check_ChildLock_Activation(void) {
            // 假设 Is_Key_Pressed(KEY_MODE) 和 Is_Key_Pressed(KEY_PLUS) 用于检测按键是否按下
            if (Is_Key_Pressed(KEY_MODE) && Is_Key_Pressed(KEY_PLUS)) {
                child_lock_press_time++;
                if (child_lock_press_time >= 300) { // 假设每10ms调用一次，3秒 = 300次
                    child_lock_press_time = 0; // 防止重复触发
                    if (g_child_lock_state == CHILD_LOCK_OFF) {
                        g_child_lock_state = CHILD_LOCK_ON;
                        // LED_Set_ChildLock_Indicator(LED_ON); // 控制LED指示
                        printf("童锁已激活\n");
                    } else {
                        g_child_lock_state = CHILD_LOCK_OFF;
                        // LED_Set_ChildLock_Indicator(LED_OFF);
                        printf("童锁已解除\n");
                    }
                }
            } else {
                child_lock_press_time = 0;
            }
        }
        ```
        """
        logger.info("模拟LLM API调用成功，返回模拟响应。")
        return simulated_response

    def _parse_llm_output(self, llm_response_text):
        """
        解析LLM返回的文本，提取结构化的代码变更指令。
        这需要根据提示词模板中定义的输出格式进行。

        参数:
            llm_response_text (str): LLM返回的原始文本。

        返回:
            list: 包含结构化变更指令的列表。每个指令是一个字典，例如：
                  `{"type": "add/replace/delete", "file_type": ".h/.c", "file_name": "...", 
                    "code_content": "...", "target_location": "...", "old_code": "..."}`
                  返回空列表如果解析失败。
        """
        logger.info("开始解析LLM的输出...")
        if not llm_response_text:
            logger.warning("LLM响应文本为空，无法解析。")
            return []

        operations = []
        # 使用正则表达式从LLM响应中提取代码块和文件名
        # 正则表达式需要根据提示词模板中约定的输出格式来设计
        # 示例：匹配 .h 和 .c 文件块
        # 匹配 `.h` 文件内容
        h_file_match = re.search(r"建议的 \.h 文件修改/内容.*?文件名: `([^`]+?)`\s*```c\n(.*?)\n```", llm_response_text, re.DOTALL)
        if h_file_match:
            file_name = h_file_match.group(1).strip()
            code_content = h_file_match.group(2).strip()
            operations.append({
                "type": "upsert", # 表示新增或完整替换文件内容 (简化处理)
                "file_type": ".h",
                "file_name": file_name,
                "code_content": code_content,
                "description": f"LLM建议对 {file_name} 的内容进行更新/创建。"
            })
            logger.info(f"解析到对 .h 文件 {file_name} 的操作。")

        # 匹配 `.c` 文件内容
        c_file_match = re.search(r"建议的 \.c 文件修改/内容.*?文件名: `([^`]+?)`\s*```c\n(.*?)\n```", llm_response_text, re.DOTALL)
        if c_file_match:
            file_name = c_file_match.group(1).strip()
            code_content = c_file_match.group(2).strip()
            operations.append({
                "type": "upsert",
                "file_type": ".c",
                "file_name": file_name,
                "code_content": code_content,
                "description": f"LLM建议对 {file_name} 的内容进行更新/创建。"
            })
            logger.info(f"解析到对 .c 文件 {file_name} 的操作。")
        
        # 针对“修改功能”的更细致解析 (如果LLM输出格式支持)
        # 例如，解析出目标文件、定位信息、旧代码、新代码
        # for modification_block in re.finditer(r"修改点 \d+:.*?目标文件: `([^`]+?)`.*?定位信息: `([^`]+?)`.*?旧代码:\s*```c\n(.*?)\n```.*?新代码:\s*```c\n(.*?)\n```", llm_response_text, re.DOTALL):
        #     file_name = modification_block.group(1).strip()
        #     location_info = modification_block.group(2).strip()
        #     old_code = modification_block.group(3).strip()
        #     new_code = modification_block.group(4).strip()
        #     file_type = ".h" if file_name.endswith(".h") else ".c"
        #     operations.append({
        #         "type": "replace_specific",
        #         "file_type": file_type,
        #         "file_name": file_name,
        #         "location_info": location_info,
        #         "old_code_snippet": old_code,
        #         "new_code_snippet": new_code,
        #         "description": f"LLM建议修改 {file_name} 中的特定代码段。"
        #     })
        #     logger.info(f"解析到对文件 {file_name} 的特定修改操作。")

        if not operations:
            logger.warning("未能从LLM响应中解析出结构化的代码操作指令。可能需要检查LLM输出格式或解析逻辑。")
            logger.debug(f"LLM原始响应内容:\n{llm_response_text}")
            # 作为回退，可以将整个LLM响应作为一种“通用建议”返回，由人工处理
            operations.append({
                "type": "general_suggestion",
                "file_type": "unknown",
                "file_name": "unknown",
                "code_content": llm_response_text, # 整个响应
                "description": "LLM提供了通用建议，未能解析为具体文件操作。"
            })

        return operations

    def generate_code_suggestions(self, requirement_item, rag_context):
        """
        主函数：为单个需求差异项生成代码建议。

        参数:
            requirement_item (dict): 单个需求差异项。
            rag_context (dict): RAG检索模块返回的上下文信息。

        返回:
            list: 包含结构化变更指令的列表。
        """
        logger.info(f"开始为需求 	'{requirement_item.get("feature_name")}	' 生成代码建议...")
        
        full_prompt = self._construct_llm_prompt(requirement_item, rag_context)
        if not full_prompt:
            logger.error("构建LLM提示词失败。")
            return []
        
        llm_response = self._call_llm_api(full_prompt)
        if not llm_response:
            logger.error("调用LLM API失败或未收到响应。")
            return []
            
        parsed_suggestions = self._parse_llm_output(llm_response)
        logger.info(f"为需求 	'{requirement_item.get("feature_name")}	' 生成了 {len(parsed_suggestions)} 条代码操作建议。")
        return parsed_suggestions

if __name__ == "__main__":
    project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
    template_file = os.path.join(project_root, "templates", "prompt_template_code_generation.md")

    # 确保模板文件存在 (通常由项目初始化阶段放置)
    if not os.path.exists(template_file):
        logger.error(f"提示词模板文件未找到: {template_file}。请先创建它。")
        # 创建一个非常简单的模拟模板用于测试
        os.makedirs(os.path.dirname(template_file), exist_ok=True)
        with open(template_file, "w", encoding="utf-8") as f_template:
            f_template.write("""
            # 嵌入式C代码生成/修改任务提示词模板 (简化版)
            任务类型: - [ ] 新增功能
            新需求描述: [此处插入新需求描述]
            旧需求描述: [此处插入相关旧需求描述]
            RAG文档: 以下是从《需求联系代码的拆解文档》中检索到的
            RAG代码: 以下是从旧项目 `TM56F1522-V1.0-OLD` 代码库中检索到的
            请生成代码。
            """)
        logger.info(f"已创建简化的模拟模板: {template_file}")

    generator = CodeGenerator(prompt_template_path=template_file)

    # 模拟输入
    mock_req_item = {
        "type": "新增",
        "feature_name": "童锁保护",
        "new_description": "长按模式键和加键3秒激活童锁，LED指示。",
        "old_description": "旧版无童锁功能。",
        "details": "防止儿童误操作。"
    }
    mock_rag_context = {
        "related_docs": [
            {"title": "童锁设计参考", "distance": 0.2, "content": "童锁应考虑按键扫描和LED状态。"}
        ],
        "related_code": [
            {
                "file_path": "key.c", "type": "function", "name": "KEY_Scan", 
                "start_line": 10, "end_line": 20, "code": "void KEY_Scan() { /* ... */ }",
                "comment": "按键扫描函数"
            },
            {
                "file_path": "led.c", "type": "function", "name": "LED_SetState", 
                "start_line": 5, "end_line": 15, "code": "void LED_SetState(int led, int state) { /* ... */ }",
                "comment": "LED设置函数"
            }
        ]
    }

    logger.info("\n--- 测试代码生成模块 ---")
    suggestions = generator.generate_code_suggestions(mock_req_item, mock_rag_context)

    if suggestions:
        logger.info("代码生成建议测试成功。生成的指令：")
        for i, op in enumerate(suggestions):
            logger.info(f"  指令 {i+1}: 类型={op.get("type")}, 文件类型={op.get("file_type")}, 文件名={op.get("file_name")}")
            logger.info(f"    描述: {op.get("description")}")
            if "code_content" in op:
                 logger.info(f"    代码 (部分): {op["code_content"][:100].replace("\n", " ")}...")
    else:
        logger.error("代码生成建议测试失败或未生成指令。")
    
    logger.info("\n代码生成器模块独立测试完成。")

